import json
import os

import httpx
import requests
from application.settings import PROMPT_PATH, ZHIPU_API_URL, ZHUPU_API_KEY


class BaseAiAgent:
    prompt_path = None  # prompt的json结构存储路径

    def init_prompt(self):
        """
        初始化prompt
        :return:
        """
        result_prompt = {}
        if self.prompt_path is None:
            raise Exception("prompt_path is None")

        path = os.path.join(PROMPT_PATH, self.prompt_path)
        file_path = os.path.join(path, "prompt.json")
        with open(file_path, "r", encoding="utf-8") as f:
            prompt_json = json.load(f)
        result_prompt['prompt_system'] = prompt_json["system"]
        result_prompt['prompt_user'] = prompt_json["user"]
        result_prompt['prompt_model'] = prompt_json["model"]
        return result_prompt

    def process(self, input_text: dict, prompt: dict):
        """
        构造请求并发送到模型接口。
        将 prompt_user 中的占位符 {key} 使用 input_text 的对应值替换。
        若缺少字段则替换为空字符串。

        :param input_text: dict 输入数据，例如 {"user_query": "...", "keywords": "...", "document": {...}}
        :param prompt: dict 由 init_prompt() 返回的 prompt 结构
        :return: dict，包含模型返回内容与 token 使用信息
        """

        # ========== 动态填充 prompt_user ==========
        user_template = prompt.get("prompt_user", "")
        # 使用 str.format_map 安全替换（避免 KeyError）

        class SafeDict(dict):
            def __missing__(self, key):
                return ""  # 缺少的键替换为空字符串

        # 格式化时将复杂对象（如 dict）转为 JSON 字符串
        safe_input = SafeDict({
            k: json.dumps(v, ensure_ascii=False) if isinstance(
                v, (dict, list)) else v
            for k, v in input_text.items()
        })
        filled_user_prompt = user_template.format_map(safe_input)

        # ========== 构造请求载荷 ==========
        payload = {
            "model": prompt["prompt_model"],
            "messages": [
                {"role": "system", "content": prompt["prompt_system"]},
                {"role": "user", "content": filled_user_prompt}
            ]
        }

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {ZHUPU_API_KEY}"
        }

        resp = requests.post(ZHIPU_API_URL, json=payload, headers=headers)
        resp_data = resp.json()

        # ========== 解析响应 ==========
        content = resp_data["choices"][0]["message"]["content"]
        usage = resp_data.get("usage", {})

        # 若需要翻译（或清洗），调用 translate()
        content = self.translate(content)

        return {"result": True, "msg": "ok", "data": {"content": content, "usage": usage}}

    def translate(self, text):

        return text

    def run(self, input_text):
        """
        运行
        :param input_text:
        :return:
        """
        prompt = self.init_prompt()
        result = self.process(input_text, prompt)
        if not result['result']:
            return {"result": False, "message": result['message']}
        return {"result": True, "message": "ok", "data": result['data']}
